"""
超声心动图数据加载器

提供超声心动图视频数据的加载、解析和管理功能
"""

import os
import cv2
import numpy as np
import torch
from pathlib import Path
from typing import List, Dict, Tuple, Optional, Union, Any
import json
import pickle
from torch.utils.data import Dataset, DataLoader
import logging

from src.utils.file_utils import get_file_list, is_video_file, is_image_file
from src.utils.logger import get_logger


class EchoVideoLoader:
    """超声心动图视频加载器"""
    
    def __init__(self, 
                 data_dir: Union[str, Path],
                 image_size: Tuple[int, int] = (1024, 1024),
                 frame_interval: int = 1,
                 max_frames: Optional[int] = None):
        """
        初始化视频加载器
        
        Args:
            data_dir: 数据目录路径
            image_size: 图像尺寸 (height, width)
            frame_interval: 帧间隔
            max_frames: 最大帧数
        """
        self.data_dir = Path(data_dir)
        self.image_size = image_size
        self.frame_interval = frame_interval
        self.max_frames = max_frames
        self.logger = get_logger("EchoVideoLoader")
        
        # 支持的文件格式
        self.video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv'}
        self.image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
        
        # 数据索引
        self.video_files = []
        self.image_sequences = []
        self.annotations = {}
        
        self._build_index()
    
    def _build_index(self):
        """构建数据索引"""
        self.logger.info(f"构建数据索引: {self.data_dir}")
        
        # 查找视频文件
        video_files = get_file_list(self.data_dir, list(self.video_extensions), recursive=True)
        self.video_files = [f for f in video_files if is_video_file(f)]
        
        # 查找图像序列
        image_files = get_file_list(self.data_dir, list(self.image_extensions), recursive=True)
        self._group_image_sequences(image_files)
        
        # 加载标注文件
        self._load_annotations()
        
        # 注意：视频文件可能不在processed目录下，而是通过video_path字段从原始数据集加载
        if len(self.video_files) == 0:
            self.logger.debug(f"在 {self.data_dir} 下未找到视频文件（这是正常的，视频文件从原始数据集路径加载）")
        else:
            self.logger.info(f"找到 {len(self.video_files)} 个视频文件")
        self.logger.info(f"找到 {len(self.image_sequences)} 个图像序列")
    
    def _group_image_sequences(self, image_files: List[Path]):
        """将图像文件分组为序列"""
        # 按目录分组
        dir_groups = {}
        for img_file in image_files:
            dir_path = img_file.parent
            if dir_path not in dir_groups:
                dir_groups[dir_path] = []
            dir_groups[dir_path].append(img_file)
        
        # 创建序列
        for dir_path, files in dir_groups.items():
            if len(files) > 1:  # 至少2帧才认为是序列
                files.sort()  # 按文件名排序
                self.image_sequences.append({
                    'dir': dir_path,
                    'files': files,
                    'type': 'image_sequence'
                })
    
    def _load_annotations(self):
        """加载标注文件"""
        annotation_files = [
            'annotations.json',
            'labels.json',
            'masks.json'
        ]
        
        for ann_file in annotation_files:
            ann_path = self.data_dir / ann_file
            if ann_path.exists():
                try:
                    with open(ann_path, 'r', encoding='utf-8') as f:
                        self.annotations = json.load(f)
                    self.logger.info(f"加载标注文件: {ann_file}")
                    break
                except Exception as e:
                    self.logger.warning(f"加载标注文件失败: {ann_file}, {e}")
    
    def load_video(self, video_path: Union[str, Path]) -> List[np.ndarray]:
        """
        加载视频文件
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            帧列表
        """
        video_path = Path(video_path)
        if not video_path.exists():
            raise FileNotFoundError(f"视频文件不存在: {video_path}")
        
        cap = cv2.VideoCapture(str(video_path))
        if not cap.isOpened():
            raise ValueError(f"无法打开视频文件: {video_path}")
        
        frames = []
        frame_count = 0
        
        try:
            while True:
                ret, frame = cap.read()
                if not ret:
                    break
                
                # 按间隔采样
                if frame_count % self.frame_interval == 0:
                    # 调整尺寸
                    frame = cv2.resize(frame, (self.image_size[1], self.image_size[0]))
                    frames.append(frame)
                    
                    # 检查最大帧数
                    if self.max_frames and len(frames) >= self.max_frames:
                        break
                
                frame_count += 1
        
        finally:
            cap.release()
        
        self.logger.info(f"加载视频: {video_path.name}, {len(frames)} 帧")
        return frames
    
    def load_image_sequence(self, sequence_info: Dict[str, Any]) -> List[np.ndarray]:
        """
        加载图像序列
        
        Args:
            sequence_info: 序列信息
            
        Returns:
            帧列表
        """
        frames = []
        files = sequence_info['files']
        
        # 限制帧数
        if self.max_frames:
            files = files[:self.max_frames * self.frame_interval:self.frame_interval]
        else:
            files = files[::self.frame_interval]
        
        for img_file in files:
            try:
                frame = cv2.imread(str(img_file))
                if frame is not None:
                    # 调整尺寸
                    frame = cv2.resize(frame, (self.image_size[1], self.image_size[0]))
                    frames.append(frame)
            except Exception as e:
                self.logger.warning(f"加载图像失败: {img_file}, {e}")
        
        self.logger.info(f"加载图像序列: {sequence_info['dir'].name}, {len(frames)} 帧")
        return frames
    
    def get_annotation(self, video_id: str) -> Optional[Dict[str, Any]]:
        """
        获取视频标注
        
        Args:
            video_id: 视频ID
            
        Returns:
            标注信息
        """
        return self.annotations.get(video_id)
    
    def list_videos(self) -> List[Path]:
        """获取所有视频文件列表"""
        return self.video_files.copy()
    
    def list_sequences(self) -> List[Dict[str, Any]]:
        """获取所有图像序列列表"""
        return self.image_sequences.copy()
    
    def get_video_info(self, video_path: Union[str, Path]) -> Dict[str, Any]:
        """
        获取视频信息
        
        Args:
            video_path: 视频文件路径
            
        Returns:
            视频信息
        """
        video_path = Path(video_path)
        cap = cv2.VideoCapture(str(video_path))
        
        if not cap.isOpened():
            return {}
        
        info = {
            'path': str(video_path),
            'fps': cap.get(cv2.CAP_PROP_FPS),
            'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)),
            'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
            'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)),
            'duration': cap.get(cv2.CAP_PROP_FRAME_COUNT) / cap.get(cv2.CAP_PROP_FPS)
        }
        
        cap.release()
        return info


class EchoDataset(Dataset):
    """超声心动图数据集"""
    
    def __init__(self, 
                 data_dir: Union[str, Path],
                 split: str = 'train',
                 image_size: Tuple[int, int] = (1024, 1024),
                 frame_interval: int = 1,
                 max_frames: Optional[int] = None,
                 load_masks: bool = True):
        """
        初始化数据集
        
        Args:
            data_dir: 数据目录
            split: 数据划分 (train/val/test)
            image_size: 图像尺寸
            frame_interval: 帧间隔
            max_frames: 最大帧数
            load_masks: 是否加载掩码
        """
        self.data_dir = Path(data_dir)
        self.split = split
        self.image_size = image_size
        self.frame_interval = frame_interval
        self.max_frames = max_frames
        self.load_masks = load_masks
        self.logger = get_logger("EchoDataset")
        
        # 初始化加载器
        self.loader = EchoVideoLoader(
            data_dir=data_dir,
            image_size=image_size,
            frame_interval=frame_interval,
            max_frames=max_frames
        )
        
        # 加载数据划分
        self.samples = self._load_split()
        
        self.logger.info(f"加载 {split} 数据集: {len(self.samples)} 个样本")
    
    def _load_split(self) -> List[Dict[str, Any]]:
        """加载数据划分"""
        split_file = self.data_dir / f"{self.split}_split.json"
        
        if split_file.exists():
            with open(split_file, 'r', encoding='utf-8') as f:
                split_data = json.load(f)
            return split_data.get('samples', [])
        else:
            # 如果没有划分文件，使用所有数据
            self.logger.warning(f"未找到划分文件: {split_file}")
            return self._create_default_split()
    
    def _create_default_split(self) -> List[Dict[str, Any]]:
        """创建默认数据划分"""
        samples = []
        
        # 添加视频样本
        for video_path in self.loader.list_videos():
            samples.append({
                'type': 'video',
                'path': str(video_path),
                'id': video_path.stem
            })
        
        # 添加图像序列样本
        for seq_info in self.loader.list_sequences():
            samples.append({
                'type': 'sequence',
                'path': str(seq_info['dir']),
                'files': [str(f) for f in seq_info['files']],
                'id': seq_info['dir'].name
            })
        
        return samples
    
    def __len__(self) -> int:
        """返回数据集大小"""
        return len(self.samples)
    
    def __getitem__(self, idx: int) -> Dict[str, Any]:
        """
        获取数据样本
        
        Args:
            idx: 样本索引
            
        Returns:
            数据样本
        """
        sample = self.samples[idx]
        
        # 加载视频/序列
        if sample['type'] == 'video':
            frames = self.loader.load_video(sample['path'])
        else:  # sequence
            frames = self.loader.load_image_sequence({
                'dir': Path(sample['path']),
                'files': [Path(f) for f in sample['files']]
            })
        
        if not frames:
            raise ValueError(f"无法加载样本: {sample['id']}")
        
        # 转换为tensor
        frames_tensor = torch.stack([
            torch.from_numpy(frame).permute(2, 0, 1).float() / 255.0
            for frame in frames
        ])
        
        result = {
            'frames': frames_tensor,
            'video_id': sample['id'],
            'num_frames': len(frames)
        }
        
        # 加载掩码（如果存在）
        if self.load_masks:
            mask_path = self.data_dir / 'masks' / f"{sample['id']}.npy"
            if mask_path.exists():
                masks = np.load(mask_path)
                result['masks'] = torch.from_numpy(masks).float()
            else:
                # 创建空掩码
                result['masks'] = torch.zeros(len(frames), *self.image_size)
        
        return result


def create_dataloader(data_dir: Union[str, Path],
                     split: str = 'train',
                     batch_size: int = 4,
                     num_workers: int = 4,
                     shuffle: bool = True,
                     **kwargs) -> DataLoader:
    """
    创建数据加载器
    
    Args:
        data_dir: 数据目录
        split: 数据划分
        batch_size: 批大小
        num_workers: 工作进程数
        shuffle: 是否打乱
        **kwargs: 其他参数
        
    Returns:
        数据加载器
    """
    dataset = EchoDataset(data_dir, split=split, **kwargs)
    
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=True
    )


def test_loader():
    """测试数据加载器"""
    logger = get_logger("TestLoader")
    
    # 测试视频加载器
    data_dir = Path("data/raw/echocardiogram")
    if data_dir.exists():
        loader = EchoVideoLoader(data_dir)
        
        logger.info(f"视频文件: {len(loader.list_videos())}")
        logger.info(f"图像序列: {len(loader.list_sequences())}")
        
        # 测试加载第一个视频
        videos = loader.list_videos()
        if videos:
            frames = loader.load_video(videos[0])
            logger.info(f"加载视频帧数: {len(frames)}")
            if frames:
                logger.info(f"帧尺寸: {frames[0].shape}")
    else:
        logger.warning(f"数据目录不存在: {data_dir}")


if __name__ == "__main__":
    test_loader()
